// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <melon/container/f14_map.h>

#include <pollux/exec/container_row_serde.h>
#include <pollux/exec/top_n.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux::exec {
    TopN::TopN(
        int32_t operatorId,
        DriverCtx *driverCtx,
        const std::shared_ptr<const core::TopNNode> &topNNode)
        : Operator(
              driverCtx,
              topNNode->outputType(),
              operatorId,
              topNNode->id(),
              "TopN"),
          count_(topNNode->count()),
          data_(std::make_unique<RowContainer>(outputType_->children(), pool())),
          comparator_(
              outputType_,
              topNNode->sortingKeys(),
              topNNode->sortingOrders(),
              data_.get()),
          topRows_(comparator_),
          decodedVectors_(outputType_->children().size()) {
        const auto numColumns{outputType_->children().size()};
        const auto numSortingKeys{topNNode->sortingKeys().size()};
        sortingKeyColumns_.reserve(numSortingKeys);
        std::vector<bool> isSortingKey(numColumns);
        for (const auto &key: topNNode->sortingKeys()) {
            sortingKeyColumns_.emplace_back(exprToChannel(key.get(), outputType_));
            isSortingKey[sortingKeyColumns_.back()] = true;
        }
        if (numColumns > numSortingKeys) {
            nonKeyColumns_.reserve(numColumns - numSortingKeys);
            for (column_index_t i = 0; i < numColumns; ++i) {
                if (!isSortingKey[i]) {
                    nonKeyColumns_.emplace_back(i);
                }
            }
        }
    }

    void TopN::addInput(RowVectorPtr input) {
        for (const auto col: sortingKeyColumns_) {
            decodedVectors_[col].decode(*input->childAt(col));
        }

        const bool hasNonKeyColumn{!nonKeyColumns_.empty()};
        // Maps passed rows of 'data_' to the corresponding input row number. These
        // input rows of non-key columns are later stored into data_.
        melon::F14FastMap<void *, vector_size_t> passedRows;
        for (auto row = 0; row < input->size(); ++row) {
            char *newRow = nullptr;
            if (topRows_.size() < count_) {
                newRow = data_->newRow();
            } else {
                char *topRow = topRows_.top();

                if (!comparator_(decodedVectors_, row, topRow)) {
                    continue;
                }
                topRows_.pop();
                // Reuse the topRow's memory.
                newRow = data_->initializeRow(topRow, true /* reuse */);
            }

            data_->initializeFields(newRow);
            for (const auto col: sortingKeyColumns_) {
                data_->store(decodedVectors_[col], row, newRow, col);
            }

            topRows_.push(newRow);
            if (hasNonKeyColumn) {
                passedRows[newRow] = row;
            }
        }

        if (hasNonKeyColumn && !passedRows.empty()) {
            for (const auto col: nonKeyColumns_) {
                decodedVectors_[col].decode(*input->childAt(col));
                for (const auto [dataRow, inputRow]: passedRows) {
                    data_->store(
                        decodedVectors_[col],
                        inputRow,
                        reinterpret_cast<char *>(dataRow),
                        col);
                }
            }
        }
    }

    RowVectorPtr TopN::getOutput() {
        if (finished_ || !noMoreInput_) {
            return nullptr;
        }

        const auto numRowsToReturn = std::min<vector_size_t>(
            outputBatchSize_, rows_.size() - numRowsReturned_);
        POLLUX_CHECK_GT(numRowsToReturn, 0);

        auto result = BaseVector::create<RowVector>(
            outputType_, numRowsToReturn, operatorCtx_->pool());

        for (auto i = 0; i < outputType_->size(); ++i) {
            data_->extractColumn(
                rows_.data() + numRowsReturned_,
                numRowsToReturn,
                i,
                result->childAt(i));
        }
        numRowsReturned_ += numRowsToReturn;
        finished_ = (numRowsReturned_ == rows_.size());
        return result;
    }

    void TopN::noMoreInput() {
        Operator::noMoreInput();
        if (topRows_.empty()) {
            finished_ = true;
            return;
        }
        rows_.resize(topRows_.size());
        for (auto i = rows_.size(); i > 0; --i) {
            rows_[i - 1] = topRows_.top();
            topRows_.pop();
        }

        outputBatchSize_ = outputBatchRows(data_->estimateRowSize());
    }

    bool TopN::isFinished() {
        return finished_;
    }
} // namespace kumo::pollux::exec
